import torch
from transformers import AutoTokenizer, AutoModelForQuestionAnswering

# ---- 1. 加载您微调好的模型和分词器 ----
# 确保这个路径指向您之前保存模型的文件夹
model_path = "./my_professional_qa_model"
print(f"正在从 '{model_path}' 加载模型和分词器...")
tokenizer = AutoTokenizer.from_pretrained(model_path)
model = AutoModelForQuestionAnswering.from_pretrained(model_path)
print("加载完成！")


def answer_question(question, context):
    """
    使用加载的模型来回答问题。

    Args:
        question (str): 你要问的问题。
        context (str): 可能包含答案的上下文段落。

    Returns:
        str: 从上下文中抽取的答案文本。
    """
    # ---- 2. 准备输入：将问题和上下文编码 ----
    # return_tensors='pt' 表示返回PyTorch张量
    inputs = tokenizer.encode_plus(question, context, add_special_tokens=True, return_tensors="pt")
    input_ids = inputs["input_ids"].tolist()[0]

    # ---- 3. 模型推理：获取答案的起始和结束位置 ----
    # 使用 torch.no_grad() 可以节约计算资源，因为我们只是推理，不需要计算梯度
    with torch.no_grad():
        outputs = model(**inputs)

    # 获取模型输出的、每个token作为答案起点和终点的分数（logits）
    answer_start_scores = outputs.start_logits
    answer_end_scores = outputs.end_logits

    # 找到分数最高的token作为答案的起点和终点
    answer_start = torch.argmax(answer_start_scores)
    answer_end = torch.argmax(answer_end_scores) + 1  # +1 是因为切片操作是左闭右开

    # ---- 4. 解码输出：将token ID转换回文本 ----
    # 从原始的 input_ids 中提取出答案部分的 token
    answer_tokens = inputs["input_ids"][0][answer_start:answer_end]
    # 使用分词器将 token ID 解码成可读的文本
    answer = tokenizer.decode(answer_tokens, skip_special_tokens=True)

    return answer


# ---- 5. 使用示例 ----
# 假设我们用之前训练时用过的数据来提问
# 这段 context 来自我们之前创建的 train.json 文件
context_from_db = "成都有很多美丽的景点，例如都江堰、熊猫谷"
question_to_ask = "成都景点有哪些？"

print("\n" + "=" * 50)
print(f"上下文 (Context): {context_from_db}")
print(f"问题 (Question): {question_to_ask}")

# 调用函数获取答案
predicted_answer = answer_question(question_to_ask, context_from_db)

print(f"\n模型预测的答案 (Answer): {predicted_answer}")
print("=" * 50)

# 另一个例子
context_finance = "成都市位于中国西南。"
question_finance = "成都位于哪里？"

print(f"\n上下文 (Context): {context_finance}")
print(f"问题 (Question): {question_finance}")
predicted_answer_finance = answer_question(question_finance, context_finance)
print(f"\n模型预测的答案 (Answer): {predicted_answer_finance}")
print("=" * 50)